from typing import Optional, List

from peewee import Model, CharField, IntegerField, FloatField, BooleanField
from playhouse.shortcuts import model_to_dict
from pydantic import BaseModel

from db.my_sqlite import DB

# 定义文件模型
class Knowledge(Model):
    id = CharField(primary_key=True, help_text="知识库id")
    userid = IntegerField(help_text="用户id")
    rag_name = CharField(help_text="知识库名称")
    description = CharField(help_text="知识库描述", null=True)
    rag_type = IntegerField(help_text="知识库类型", null=True)
    embed = CharField(help_text="embedding模型名称", null=True)
    llm = CharField(help_text="llm模型名称", null=True)
    reranker = CharField(help_text="reranker模型名称", null=True)
    chunk_type = IntegerField(help_text="分片类型", null=True)
    chunk_size = IntegerField(help_text="分片大小", null=True)
    chunk_overlap = IntegerField(help_text="分片重叠长度", null=True)
    content_augmentation = IntegerField(help_text="上下文增强类型", null=True)
    retrieval_type = IntegerField(help_text="检索类型", null=True)
    mix_threshold = FloatField(help_text="混合检索阈值", null=True)
    top_k = IntegerField(help_text="获取重排后结果数量", null=True)
    is_set = BooleanField(help_text="参数是否配置完成", null=True)
    create_date= CharField(help_text="创建时间", null=True)
    update_date = CharField(help_text="修改时间", null=True)

    class Meta:
        database = DB


class KnowledgeModel(BaseModel):
    id: Optional[str]
    userid: Optional[int]
    rag_name: Optional[str]
    description: Optional[str]
    embed: Optional[str]
    llm: Optional[str]
    reranker: Optional[str]
    chunk_size: Optional[int]
    chunk_overlap: Optional[int]
    chunk_type: Optional[int] = 1
    rag_type: Optional[int] = 1
    content_augmentation: Optional[int] = 2
    retrieval_type: Optional[int] = 3
    mix_threshold: Optional[float] = 0.75
    top_k: Optional[int] = 5
    is_set: Optional[bool] = False
    create_date:Optional[str]
    update_date:Optional[str]
    query: Optional[str] = None
class KnowledgeTable:
    def __init__(self, db):
        self.db = db
        self.db.create_tables([Knowledge])

    @staticmethod
    def delete_knowledge(userid: int, id: str) -> bool:
        return Knowledge.delete().where(
            (Knowledge.userid == userid) &
            (Knowledge.id == id)
        ).execute() > 0

    @staticmethod
    def update_knowledge(userid: int, id: str, old_kg: KnowledgeModel, new_kg: KnowledgeModel) -> bool:
        return Knowledge.update(
            rag_name=new_kg.rag_name if new_kg.rag_name else old_kg.rag_name,
            description=new_kg.description if new_kg.description else old_kg.description,
            rag_type=new_kg.rag_type if new_kg.rag_type else old_kg.rag_type,
            embed=new_kg.embed if new_kg.embed else old_kg.embed,
            llm=new_kg.llm if new_kg.llm else old_kg.llm,
            reranker=new_kg.reranker if new_kg.reranker else old_kg.reranker,
            chunk_type=new_kg.chunk_type if new_kg.chunk_type else old_kg.chunk_type,
            chunk_size=new_kg.chunk_size if new_kg.chunk_size else old_kg.chunk_size,
            chunk_overlap=new_kg.chunk_overlap if new_kg.chunk_overlap else old_kg.chunk_overlap,
            content_augmentation=new_kg.content_augmentation if new_kg.content_augmentation else old_kg.content_augmentation,
            retrieval_type=new_kg.retrieval_type if new_kg.retrieval_type else old_kg.retrieval_type,
            mix_threshold=new_kg.mix_threshold if new_kg.mix_threshold else old_kg.mix_threshold,
            top_k=new_kg.top_k if new_kg.top_k else old_kg.top_k,
            is_set=True,
            update_date = new_kg.update_date
        ).where(
            (Knowledge.userid == userid) &
            (Knowledge.id == id)
        ).execute() > 0

    @staticmethod
    def get_knowledge(userid: int, id: str) -> Optional[KnowledgeModel]:
        file = Knowledge.get_or_none(
            (Knowledge.userid == userid) &
            (Knowledge.id == id)
        )
        if file:
            return KnowledgeModel(**model_to_dict(file))
        return None

    @staticmethod
    def get_all_knowledge(query:str,userid: int) -> List[KnowledgeModel]:
        base_query = Knowledge.select().where(
            (Knowledge.userid == userid)
        )
        if query:
            base_query = base_query.where(Knowledge.rag_name ** f"%{query}%")  # 添加模糊查询条件

        return [KnowledgeModel(**model_to_dict(file)) for file in base_query]

    @staticmethod
    def create_knowledge(userid: int, id: str, rag_name: str, description: str,create_date:str) -> Optional[KnowledgeModel]:
        file = Knowledge.create(
            id=id,
            userid=userid,
            rag_name=rag_name,
            description=description,
            create_date=create_date,
            update_date=create_date
        )
        return True

    @staticmethod
    def get_public_knowledge(userid: int) -> List[KnowledgeModel]:
        files = Knowledge.select().where(
            (Knowledge.userid == userid)
        )
        return [KnowledgeModel(**model_to_dict(file)) for file in files]

KG = KnowledgeTable(DB)
